#include "llvm/Pass.h"
#include "llvm/Module.h"
#include "llvm/Function.h"
#include "llvm/Instructions.h"
#include "llvm/CallingConv.h"
#include "llvm/DerivedTypes.h"
#include "llvm/InstrTypes.h"
#include "llvm/Constants.h"
#include "llvm/Support/CFG.h"
#include "llvm/Target/TargetData.h"
#include <vector>
#include <set>
#include <list>
#include <map>
#include <iostream>
#include <sstream>
#include <llvm/ADT/STLExtras.h>
#include <llvm/Target/TargetMachine.h>

using namespace llvm;
using namespace std;

namespace {

    class ArrayCheckPass : public FunctionPass {
        //Call to array bound check function defined STATICALLY in another class
        Constant *arr_chk_call;
        Constant *debug_call;
        Constant *dyn_chk_call;
        std::map<GetElementPtrInst *, Value *> hmpIndValueDefs;
        std::map<GetElementPtrInst *, Instruction *> hmpIndInst;
        std::map<Value *, int > hmpUniqChecksPerValue;
        std::map<Value *, std::list<GetElementPtrInst *> > hmpDynamicArrayRefs;
        std::map<GetElementPtrInst *, GetElementPtrInst *> hmpDArrayChains;
        std::map<Value *, GetElementPtrInst *> hmpLastElems;
        std::set<GetElementPtrInst *> setConveredAccesses;
        std::set<BasicBlock *> setMarkedBBs;

    public:
        static char ID;

        ArrayCheckPass() : FunctionPass(&ID) {
        }

        virtual bool doInitialization(Module &mdl) {

            //Get the elemLinked check_bounds function
            arr_chk_call = mdl.getOrInsertFunction("check_bounds",
                    /* return type */ Type::getVoidTy(mdl.getContext()),
                    /* value */ Type::getInt32Ty(mdl.getContext()),
                    /* lim */ Type::getInt32Ty(mdl.getContext()),
                    NULL);
            if (!arr_chk_call) {
                std::cerr << "### Error: fcall is NULL\n" << std::endl;
            }

            // Set the calling convention to C, so we interoperate properly with C code.
            {
                Function *tmp = cast<Function > (arr_chk_call);
                tmp->setCallingConv(CallingConv::C);
            }

            debug_call = mdl.getOrInsertFunction("debug_func",
                    /* return type */ Type::getVoidTy(mdl.getContext()),
                    /* value */ Type::getInt32Ty(mdl.getContext()),
                    NULL);
            if (!debug_call) {
                std::cerr << "### Error: debug_call is NULL\n" << std::endl;
            }

            // Set the calling convention to C, so we interoperate properly with C code.
            {
                Function *tmp = cast<Function > (debug_call);
                tmp->setCallingConv(CallingConv::C);
            }

            //int value, int baseLim, int divisor
            dyn_chk_call = mdl.getOrInsertFunction("dyn_check_bounds",
                    /* return type */ Type::getInt32Ty(mdl.getContext()),
                    /* value */ Type::getInt32Ty(mdl.getContext()),
                    /* baseLim */ Type::getInt32Ty(mdl.getContext()),
                    /* divisor */ Type::getInt32Ty(mdl.getContext()),
                    NULL);
            if (!dyn_chk_call) {
                std::cerr << "### Error: calc_size_call is NULL\n" << std::endl;
            }

            // Set the calling convention to C, so we interoperate properly with C code.
            {
                Function *tmp = cast<Function > (dyn_chk_call);
                tmp->setCallingConv(CallingConv::C);
            }
            //true, since we modified the module
            return true;
        }

        //Checks whether the given parameter value is subsumed by previous checks on the same value
        //param : parameter on which check is applied
        //iNSize : upper bound being checked on the parameter

        bool isUniqueCheckPerVal(Value *param, int iNSize) {
            //std::cerr << "Checking " << *param << "\nwith " << iNSize << "\n";
            //Whether any bound exists for this value
            bool bUnq = hmpUniqChecksPerValue.find(param) == hmpUniqChecksPerValue.end();

            //current bound value stored for this param
            int curMinForVal = hmpUniqChecksPerValue[param];

            //If no bound is present, definitely NOT subsumed
            bool bSubsume = !bUnq;
            if (bUnq) {
                //If unique, store this as the bound
                hmpUniqChecksPerValue[param] = iNSize;
            } else if (iNSize > curMinForVal) {
                //if not unique, this check is subsumed by lower bounds
                bSubsume = true;
            } else {
                //if this bound tightens the previous one on this param, update the bound
                hmpUniqChecksPerValue[param] = iNSize;
            }
            return !bSubsume;
        }

        std::map<CallInst *, LoadInst *> hmpAddedLoads;

        //Adds array checks for the allocation instruction 'inst',
        //vctSizes = vector of upper bounds  for multidimensional arrays
        //sz_index = index in the vctSizes (i.e. subscript)
        //F = Current Function (needed to create call for bounds check)
        //numBits = number of bits for single array element
        //setOfIndices = set of indices already checked for each subscript

        void addArrayChecks(Instruction *inst, vector<const ArrayType*> &vctSizes, unsigned int sz_index, Function &F, vector<set<Value *> > &setOfIndices) {
            if (sz_index >= vctSizes.size())
                return;
            //Get the current upper bound
            int iNSize = vctSizes[sz_index]->getNumElements();
            //Create parameters to be passed to check_bounds function
            //0th parameter is substituted later
            Value * params[] = {NULL,
                ConstantInt::get(Type::getInt32Ty(F.getContext()), iNSize)};
            //Get the previously checked indices
            set <Value *> &indVals = setOfIndices[sz_index];
            for (Value::use_iterator i = inst->use_begin(), e = inst->use_end(); i != e; ++i) {
                GetElementPtrInst * getElmInst = dyn_cast<GetElementPtrInst > (*i);
                if (getElmInst == NULL)
                    continue;
                int numOps = getElmInst->getNumOperands();
                //Get the accessed index
                Value *index = getElmInst->getOperand(numOps - 1);
                params[0] = index;
                //if index is NOT checked previously
                if (indVals.find(params[0]) == indVals.end()) {
                    //check if this index is an alias of other index
                    if (hmpIndValueDefs.find(getElmInst) != hmpIndValueDefs.end()) {
                        //Get the stored value for the index
                        Value *store_value = hmpIndValueDefs[getElmInst];
                        //Get the instruction at which index value was stored
                        //We will add the check immediately after this instruction
                        //Thus this check reaches EVERY use
                        Instruction *chk_inst = hmpIndInst[getElmInst];
                        if (isUniqueCheckPerVal(store_value, iNSize)) {
                            params[0] = store_value;
                            if (chk_inst != NULL) {
                                //Create iterator at the instruction
                                BasicBlock::iterator next(chk_inst);
                                ++next; // After this line, it refers to the instruction after *inst
                                //Check for end of basic block
                                //Create and call check_bounds
                                if (next != chk_inst->getParent()->end()) {
                                    CallInst::Create(arr_chk_call, params, array_endof(params), "", next);
                                } else {
                                    CallInst::Create(arr_chk_call, params, array_endof(params), "", chk_inst->getParent());
                                }
                            }
                        } else {
                            //not unique CHECK !!
                            //std::cerr << "DUPLICATE/SUBSUMED CHECK !!!" << "\n";
                        }
                    } else {
                        //Check for uniqueness
                        if (isUniqueCheckPerVal(params[0], iNSize)) {
                            //Create and call check_bounds
                            CallInst::Create(arr_chk_call, params, array_endof(params), "", getElmInst);
                        } else {
                            //not unique CHECK !!
                            //std::cerr << "DUPLICATE/SUBSUMED CHECK !!!" << "\n";
                        }
                    }
                    //Mark this index value as checked
                    indVals.insert(params[0]);
                } else {
                    //not unique CHECK !!
                    //std::cerr << "DUPLICATE/SUBSUMED CHECK !!!" << "\n";
                }
                setConveredAccesses.insert(getElmInst);
                //Update the used index values
                setOfIndices[sz_index] = indVals;
                if (numOps > 2) {
                    //if multidimensional array, recursive call
                    addArrayChecks((Instruction *) getElmInst, vctSizes, sz_index + 1, F, setOfIndices);
                }
            }
        }

        void searchBBForAlias(BasicBlock *bb, GetElementPtrInst* &getElemInst, Value* &defValue, Instruction* & defInst, LoadInst* &ldInst) {
            if (setMarkedBBs.find(bb) != setMarkedBBs.end())
                return;
            setMarkedBBs.insert(bb);
            //Start from the end of basic block
            BasicBlock::iterator rev_itr = bb->end();
            //For the current basic block (i.e. same as load instruction), start from load instruction's previous instruction
            if (ldInst->getParent() == bb) {
                rev_itr = BasicBlock::iterator(ldInst);
            }
            if (rev_itr != bb->begin()) {
                //Traverse back
                do {
                    rev_itr--;
                    Value *operand = NULL, *curValue = NULL;
                    Instruction *curDef = NULL;
                    if (CallInst::classof(rev_itr)) {
                        //If it is call instruction, return value (instruction itself in LLVM) is the alias
                        operand = dyn_cast<Value > (rev_itr);
                        curValue = operand;
                        //Defining instruction is the call instruction
                    } else {
                        //If it is store instruction, 0th operand value is the alias
                        StoreInst *storInst = dyn_cast<StoreInst > (rev_itr);
                        if (storInst == NULL)
                            continue;
                        operand = storInst->getPointerOperand();
                        curValue = storInst->getOperand(0);
                    }
                    if (operand == NULL)
                        continue;
                    if (operand != ldInst->getPointerOperand()) {
                        if (curValue == operand)
                            break;
                    }
                    //Defining instriction is the current Instruction (store / call)
                    curDef = rev_itr;
                    if (defValue == NULL) {
                        //If no previous definition is found, set current value as the defn,
                        //and current defInst as the defining instruction
                        defValue = curValue;
                        defInst = curDef;
                    } else if (defValue != curValue) {
                        //otherwise, two definitions are reaching the use, i.e. NOT an alias
                        //std::cerr << "!!! Found different values for " << ldInst->getPointerOperand() << "\n";
                        ldInst = NULL;
                    }
                    return;

                } while (rev_itr != bb->begin());
            }
            if (defValue == NULL) {
                //if no defining value is found, search predecessor(s)
                //std::cerr << "Examining predecessor(s) \n";
                for (pred_iterator PI = pred_begin(bb), E = pred_end(bb); PI != E; ++PI) {
                    BasicBlock *Pred = *PI;
                    //Recursively search this basic block predecessor
                    searchBBForAlias(Pred, getElemInst, defValue, defInst, ldInst);
                    if (ldInst == NULL) {
                        //If two reaching definitions are found,NO alias
                        defValue = NULL;
                        return;
                    }
                    if (defValue == NULL) {
                        //If no definition was found, NO alias
                        ldInst = NULL;
                        return;
                    }
                }
            }
        }

        long getSize(const Type* type) {
            //std::cerr << "In getSize, type = " << type->getTypeID() << "\n";
            if (type->isSingleValueType())
                return type->getPrimitiveSizeInBits() / 8;
            long total = 0;
            const StructType *struct_typ = dyn_cast<StructType > (type);
            if (struct_typ != NULL) {
                for (unsigned i = 0, lim = struct_typ->getNumElements(); i < lim; i++) {
                    const Type *subType = struct_typ->getTypeAtIndex(i);
                    int subSize = getSize(subType);
                    total += subSize;
                }
            }
            if (ArrayType::classof(type)) {
                int subTotal = 0;
                while (ArrayType::classof(type)) {
                    const ArrayType *arr_type = dyn_cast<ArrayType > (type);
                    type = arr_type->getElementType();
                    int numElems = arr_type->getNumElements();
                    int subSize = (numElems == 0 ? 1 : numElems) * getSize(type);
                    subTotal += subSize;
                }
                total += subTotal;
            }
            return total;
        }

        void recSearchFromRef(GetElementPtrInst* &parent, Instruction* & current, GetElementPtrInst* &lastInst) {
            for (Value::use_iterator usr_itr = current->use_begin(), end = current->use_end(); usr_itr != end; usr_itr++) {
                Instruction *instr = dyn_cast<Instruction > (usr_itr);
                if (instr == NULL)
                    continue;
                GetElementPtrInst *arrayAccess = dyn_cast<GetElementPtrInst > (instr);
                if (arrayAccess != NULL && !StructType::classof(arrayAccess->getPointerOperandType()->getElementType())) {
                    hmpDArrayChains[parent] = arrayAccess;
                    lastInst = arrayAccess;
                    recSearchFromRef(arrayAccess, instr, lastInst);
                } else
                    recSearchFromRef(parent, instr, lastInst);
            }
        }

        void recSearchArrayRef(Value * &allocInstruction, Value* &instr) {
            for (Value::use_iterator usr_itr = instr->use_begin(), end = instr->use_end(); usr_itr != end; usr_itr++) {
                Value *usr_val = *usr_itr;
                GetElementPtrInst *arrayAccess = dyn_cast<GetElementPtrInst > (usr_itr);
                if (arrayAccess != NULL) {
                    //std::cerr << "Found : " << *arrayAccess << "\n";
                    hmpDynamicArrayRefs[allocInstruction].push_back(arrayAccess);
                    Instruction* inst = (Instruction *) arrayAccess;
                    GetElementPtrInst* last = arrayAccess;
                    recSearchFromRef(arrayAccess, inst, last);
                    hmpLastElems[allocInstruction] = last;
                } else
                    recSearchArrayRef(allocInstruction, usr_val);
            }
        }

        virtual bool runOnFunction(Function &F) {
            if (!F.getNameStr().compare("check_bounds"))
                return false;
            if (!F.getNameStr().compare("dyn_check_bounds"))
                return false;
            if (!F.getNameStr().compare("debug_func"))
                return false;
            //std::cerr << "Inside function " << F.getNameStr() << " : " << F.getParent()->getModuleIdentifier() << "\n";
            bool DEBUG = false;
            if (DEBUG) {
                for (Function::iterator bb = F.begin(), bb_end = F.end(); bb != bb_end; bb++) {
                    for (BasicBlock::iterator instr = bb->begin(), instr_end = bb->end(); instr != instr_end; instr++) {
                        std::cerr << *instr << "\n";
                    }
                }
                return false;
            }
            //clear the collections
            hmpAddedLoads.clear();
            hmpDArrayChains.clear();
            hmpDynamicArrayRefs.clear();
            hmpIndInst.clear();
            hmpIndValueDefs.clear();
            hmpLastElems.clear();
            hmpUniqChecksPerValue.clear();
            setConveredAccesses.clear();

            for (Function::iterator bb = F.begin(), bb_end = F.end(); bb != bb_end; bb++) {
                for (BasicBlock::iterator instr = bb->begin(), instr_end = bb->end(); instr != instr_end; instr++) {
                    GetElementPtrInst *getElemInst = dyn_cast<GetElementPtrInst > (instr);
                    if (getElemInst == NULL)
                        continue;
                    //if the getelementpt instruction (i.e. array access ) uses load instruction,
                    //we invoke alias analysis
                    LoadInst * ldInst = dyn_cast<LoadInst > (getElemInst->getOperand(getElemInst->getNumOperands() - 1));
                    if (ldInst != NULL) {
                        //std::cerr << "Load : " << *ldInst << "\n";
                        Value *defValue = NULL;
                        Instruction *defInst = NULL;
                        //Starting from the current basic block, search for defining alias
                        setMarkedBBs.clear();
                        searchBBForAlias(bb, getElemInst, defValue, defInst, ldInst);
                        if (ldInst != NULL && defValue != NULL) {
                            //if found, store it in the hashmap for this Array Access Instruction
                            hmpIndValueDefs[getElemInst] = defValue;
                            hmpIndInst[getElemInst] = defInst;
                            //std::cerr << *defInst << "\n";
                        }
                    }
                }
            }
            for (Function::iterator bb = F.begin(), bb_end = F.end(); bb != bb_end; bb++) {
                //hmpAddedLoads.clear();
                for (BasicBlock::iterator instr = bb->begin(), instr_end = bb->end(); instr != instr_end; instr++) {
                    //Check if this is an allocation instruction
                    AllocationInst *alloc_inst = dyn_cast<AllocationInst > (instr);
                    if (alloc_inst == NULL)
                        continue;
                    //std::cerr << *alloc_inst << "\n";
                    if (MallocInst::classof(alloc_inst))
                        continue;
                    const Type *type = alloc_inst->getAllocatedType();
                    //Check whether this is dynamic array allocation
                    if (alloc_inst->isArrayAllocation()) {
                        //std::cerr << *alloc_inst << "\n";
                        Value *totalArraySize = alloc_inst->getArraySize();
                        BitCastInst *bitCast = dyn_cast<BitCastInst > (alloc_inst->use_begin());
                        if (bitCast == NULL)
                            continue;
                        StoreInst *storInst = dyn_cast<StoreInst > (bitCast->use_begin());
                        if (storInst == NULL)
                            continue;
                        Value *arrayDef = storInst->getOperand(1);
                        //std::cerr << "Used in : \n";
                        recSearchArrayRef(arrayDef, arrayDef);
                        GetElementPtrInst *sizeInst = hmpLastElems[arrayDef];
                        //std::cerr << "Size = " << *sizeInst << "\n";
                        //Skip unused defs
                        if (sizeInst == NULL)
                            continue;
                        int elem_size = getSize(sizeInst->getPointerOperand()->getType()->getContainedType(0));
                        //std::cerr << "elem_size = " << elem_size << "\n";
                        Value *type_size = ConstantInt::get(Type::getInt32Ty(F.getContext()), elem_size);

                        std::list<GetElementPtrInst*> &elmChain = hmpDynamicArrayRefs[arrayDef];
                        //std::cerr << "locations = " << elmChain.size() << "\n";
                        std::map<GetElementPtrInst *, Value *> hmpFwdLims;
                        std::map<GetElementPtrInst *, Value *> hmpIndices;
                        for (std::list<GetElementPtrInst*>::iterator l_itr = elmChain.begin(), end = elmChain.end(); l_itr != end; l_itr++) {
                            GetElementPtrInst *start = *l_itr;
                            while (true) {
                                Value *actualIndex = start->getOperand(start->getNumOperands() - 1);
                                BinaryOperator *exIndex = dyn_cast<BinaryOperator > (actualIndex);
                                if (exIndex != NULL) {
                                    actualIndex = exIndex->getOperand(0);
                                    hmpFwdLims[start] = exIndex->getOperand(1);
                                } else {
                                    hmpFwdLims[start] = ConstantInt::get(Type::getInt32Ty(F.getContext()), getSize(start->getOperand(0)->getType()->getContainedType(0)));
                                }
                                hmpIndices[start] = actualIndex;
                                //special case for last index
                                if (hmpDArrayChains.find(start) == hmpDArrayChains.end()) {
                                    //std::cerr << "Modifying size" << "\n";
                                    hmpFwdLims[start] = type_size;
                                    break;
                                }
                                //std::cerr << "Index is " << *actualIndex << "\n";
                                //std::cerr << "Lim is " << *hmpFwdLims[start] << "\n";
                                start = hmpDArrayChains[start];
                            }
                            start = *l_itr;
                            //add checks for indices EXCEPT first and last
                            int chainSize = 0;
                            while (hmpDArrayChains.find(start) != hmpDArrayChains.end()) {
                                GetElementPtrInst *elemLink = hmpDArrayChains[start];
                                if (hmpFwdLims.find(elemLink) == hmpFwdLims.end())
                                    break;
                                Value * params[] = {
                                    hmpIndices[elemLink],
                                    hmpFwdLims[start],
                                    hmpFwdLims[elemLink]
                                };
                                CallInst::Create(dyn_chk_call, params, array_endof(params), "", elemLink);
                                setConveredAccesses.insert(elemLink);
                                chainSize++;
                                start = elemLink;
                            }
                            //Special case for first index
                            {
                                start = *l_itr;
                                Value * params[] = {
                                    hmpIndices[start],
                                    totalArraySize,
                                    hmpFwdLims[start]
                                };
                                CallInst::Create(dyn_chk_call, params, array_endof(params), "", start);
                                setConveredAccesses.insert(start);
                            }
                        }
                    } else {
                        //Statically allocated arrays
                        const ArrayType *arrType = dyn_cast<ArrayType > (type);
                        if (arrType == NULL)
                            continue;
                        vector<const ArrayType*> vctSubArrays;
                        vector<set<Value *> > vctIndices;
                        const ArrayType * tr = arrType;
                        while (true) {
                            //Create vector of sizes for all dimension(s)
                            vctSubArrays.push_back(tr);
                            set<Value *> vals;
                            vctIndices.push_back(vals);
                            const ArrayType *subType = dyn_cast<ArrayType > (tr->getElementType());
                            if (subType == NULL)
                                break;
                            //std::cerr << subType->getNumElements() << "\n";
                            tr = subType;
                        }
                        //Add array checks recursively
                        addArrayChecks((Instruction *) alloc_inst, vctSubArrays, 0, F, vctIndices);
                    }
                }
            }

            //Go for arrays inside structures
            for (Function::iterator bb = F.begin(), bb_end = F.end(); bb != bb_end; bb++) {
                for (BasicBlock::iterator instr = bb->begin(), instr_end = bb->end(); instr != instr_end; instr++) {
                    //Check if this is an allocation instruction
                    GetElementPtrInst *getElemInst = dyn_cast<GetElementPtrInst > (instr);
                    //Skip non-array access instruction
                    if (getElemInst == NULL)
                        continue;
                    //skip instructions already covered
                    if (setConveredAccesses.find(getElemInst) != setConveredAccesses.end())
                        continue;
                    const ArrayType *arr_type = dyn_cast<ArrayType > (getElemInst->getPointerOperandType()->getElementType());
                    //skip structure access instructions (structures are also accessed by getelementptr !!)
                    if (arr_type == NULL)
                        continue;
                    Value *idx = getElemInst->getOperand(getElemInst->getNumOperands() - 1);
                    int iNSize = arr_type->getNumElements();
                    if (isUniqueCheckPerVal(idx, iNSize)) {
                        Value * params[] = {idx, ConstantInt::get(Type::getInt32Ty(F.getContext()), iNSize)};
                        CallInst::Create(arr_chk_call, params, array_endof(params), "", getElemInst);
                    }
                }
            }
            return true;
        }
    };
}

char ArrayCheckPass::ID = 0;
static RegisterPass<ArrayCheckPass> X("archkpass", "Array Check pass");
